package com.zytech.zspider.task.impl.http;

import com.gargoylesoftware.htmlunit.HttpMethod;
import com.gargoylesoftware.htmlunit.Page;
import com.gargoylesoftware.htmlunit.WebRequest;
import com.gargoylesoftware.htmlunit.WebResponse;
import com.gargoylesoftware.htmlunit.util.NameValuePair;
import com.google.common.collect.Maps;
import com.zytech.modules.lang.AggregateFunction;
import com.zytech.modules.lang.Function;
import com.zytech.modules.lang.IOs;
import com.zytech.modules.lang.Langs;
import com.zytech.modules.lang.annotations.Stateless;
import com.zytech.zspider.core.TaskDefinition;
import com.zytech.zspider.core.TaskDefinitionHelper;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.Map;

/**
 * @author siuming
 */
@Stateless
class WebClientTaskRequestExecutor implements HttpTaskRequestExecutor {

    @Override
    public HttpTaskResponse submit(HttpTaskRequest request, Map<String, String> settings) {

        InputStream in = null;
        try {

            WebRequest webRequest = transformWebRequest(request, settings);
            Page page = WebClientFactory.of(settings).getPage(webRequest);
            WebResponse response = page.getWebResponse();
            Map<String, String> headers = Langs.aggregate(Maps.<String, String>newHashMap(), response.getResponseHeaders(), new AggregateFunction<Map<String, String>, NameValuePair>() {
                @Override
                public Map<String, String> apply(Map<String, String> result, NameValuePair pair) {
                    result.put(pair.getName(), pair.getValue());
                    return result;
                }
            });

            in = response.getContentAsStream();
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            IOs.piping(in, out);

            return HttpTaskResponse
                    .ofBuilder()
                    .headers(headers)
                    .content(out.toByteArray())
                    .charset(response.getContentCharset())
                    .mimeType(response.getContentType())
                    .wrapContent(page)
                    .build();
        } catch (IOException e) {
            throw Langs.toUncheck(e);
        } finally {
            if (null != in) {
                IOs.freeQuietly(in);
            }
        }
    }

    private WebRequest transformWebRequest(HttpTaskRequest request, Map<String, String> props) {
        try {
            String charset = TaskDefinitionHelper.getSetting(props, TaskDefinition.Settings.HTTP_ENCODING, "UTF-8");

            WebRequest webRequest = new WebRequest(new URL(request.getUrl()));
            webRequest.setCharset(charset);
            webRequest.setHttpMethod(HttpMethod.valueOf(request.getVerb().name().toUpperCase()));
            webRequest.setAdditionalHeaders(Maps.newHashMap(request.getHeaders()));

            if (request.getVerb() == HttpTaskRequest.Verb.POST) {
                webRequest.setRequestBody(request.getBody());
            }

            Iterable<NameValuePair> params = Langs.map(request.getParams().entrySet(), new Function<Map.Entry<String, String>, NameValuePair>() {
                @Override
                public NameValuePair apply(Map.Entry<String, String> entry) {
                    return new NameValuePair(entry.getKey(), entry.getValue());
                }
            });
            webRequest.setRequestParameters(Langs.aslist(params));
            return webRequest;
        } catch (MalformedURLException e) {
            throw Langs.toUncheck(e);
        }
    }

    @Override
    public void shutdown() {
        // TODO
    }
}
